#!/bin/bash
#set -xe
# search_linux_definitions.sh - Linux内核定义搜索工具

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 默认搜索目录（可以修改为您的内核源码路径）
DEFAULT_DIRS=(
    "$HOME/linux-2.6.10/include"
    "$HOME/linux-2.6.10/arch/i386"
    "$HOME/linux-2.6.10/kernel"
    # "mm"
    # "fs"
    # "drivers"
    # "net"
    # "lib"
)

# 使用说明
usage() {
    echo -e "${GREEN}Linux内核定义搜索工具${NC}"
    echo "用法: $0 [选项] <关键词>"
    echo ""
    echo "选项:"
    echo "  -d, --dirs <目录列表>      指定搜索目录，用逗号分隔"
    echo "  -e, --ext <扩展名列表>     指定文件扩展名，用逗号分隔"
    echo "  -a, --arch <架构类型列表>  指定架构, 用逗号分隔"
    echo "  -t, --type <类型>         搜索类型(默认是all): macro, function, struct, enum_const, typedef, all"
    echo "  -n  --num <行数>          显示最佳匹配行的前后各几行"
    echo "  -s  --save <0/1>          是否保存结果输出文件, 默认不保存"
    echo "  -v, --verbose             显示详细输出"
    echo "  -h, --help                显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 -t macro BITS_PER_LONG"
    echo "  $0 -t function set_bit"
    echo ""
}

# 初始化变量
SEARCH_DIRS=()
FILE_EXTENSIONS=()
ARCH=()
NUM=5
SAVE=0
SEARCH_TYPE="all"
VERBOSE=0
KEYWORD=""

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -d|--dirs)
            if [[ $# -lt 2 ]]; then
                echo -e "${RED}错误: $1 选项需要参数${NC}"
                usage
                exit 1
            fi
            IFS=',' read -ra SEARCH_DIRS <<< "$2"
            shift 2
            ;;
        -e|--ext)
            if [[ $# -lt 2 ]]; then
                echo -e "${RED}错误: $1 选项需要参数${NC}"
                usage
                exit 1
            fi
            IFS=',' read -ra FILE_EXTENSIONS <<< "$2"
            shift 2
            ;;
        -a|--arch)
            if [[ $# -lt 2 ]]; then
                echo -e "${RED}错误: $1 选项需要参数${NC}"
                usage
                exit 1
            fi
            IFS=',' read -ra ARCH <<< "$2"
            shift 2
            ;;
        -n|--num)
            if [[ $# -lt 2 ]]; then
                echo -e "${RED}错误: $1 选项需要参数${NC}"
                usage
                exit 1
            fi
            NUM="$2"
            shift 2
            ;;
        -s|--save)
            if [[ $# -lt 2 ]]; then
                echo -e "${RED}错误: $1 选项需要参数${NC}"
                usage
                exit 1
            fi
            SAVE="$2"
            shift 2
            ;;
        -t|--type)
            if [[ $# -lt 2 ]]; then
                echo -e "${RED}错误: $1 选项需要参数${NC}"
                usage
                exit 1
            fi
            SEARCH_TYPE="$2"
            # 定义合法类型数组
            valid_types=("macro" "function" "struct" "enum_const" "typedef" "all")
            # 验证参数是否合法
            if [[ ! " ${valid_types[@]} " =~ " ${SEARCH_TYPE} " ]]; then
                echo -e "${RED}错误: 类型 '$SEARCH_TYPE' 无效${NC}"
                echo -e "${YELLOW}支持的类型: ${valid_types[*]}${NC}"
                usage
                exit 1
            fi
            shift 2
            ;;
        -v|--verbose)
            VERBOSE=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        -*)
            echo -e "${RED}错误: 未知选项 $1${NC}"
            usage
            exit 1
            ;;
        *)
            KEYWORD="$1"
            shift
            ;;
    esac
done

# 检查关键词是否提供
if [[ -z "$KEYWORD" ]]; then
    echo -e "${RED}错误: 必须提供搜索关键词${NC}"
    usage
    exit 1
fi

# 设置默认目录（如果未指定）
if [[ ${#SEARCH_DIRS[@]} -eq 0 ]]; then
    SEARCH_DIRS=("${DEFAULT_DIRS[@]}")
fi

# 设置默认文件扩展名（如果未指定）
if [[ ${#FILE_EXTENSIONS[@]} -eq 0 ]]; then
    FILE_EXTENSIONS=("h" "c" "S")
fi

# 设置默认架构类型（如果未指定）
if [[ ${#ARCH[@]} -eq 0 ]]; then
    ARCH=("i386")
fi

# 构建文件扩展名匹配模式
build_extension_pattern() {
    local pattern=""
    for ext in "${FILE_EXTENSIONS[@]}"; do
        if [[ -n "$pattern" ]]; then
            pattern="${pattern}\|"
        fi
        pattern="${pattern}\\.$ext\$"
    done
    echo "$pattern"
}

# 构建搜索模式
build_search_pattern() {
    local keyword="$1"
    local type="$2"
    
    local keyword_reg="$(printf '\x1b\[[0-9;]*m')$keyword$(printf '\x1b\[[0-9;]*m')"
    local not_char_reg="[^a-zA-Z0-9_-]"
    local start_reg="^.*:[0-9]+:[[:blank:]]*"
    local blank_one_more_reg="[[:blank:]]+"
    local blank_zero_more_reg="[[:blank:]]*"
    case "$type" in
        macro)
            # 宏定义模式：以 #define 开头
            echo "${start_reg}#${blank_zero_more_reg}define${blank_one_more_reg}${keyword_reg}${not_char_reg}"
            ;;
        function)
            # 函数定义模式：修饰符 + inline + 返回类型 + 函数名 + 参数
            echo "${start_reg}[a-zA-Z_][^\(]*${blank_one_more_reg}${keyword_reg}${blank_zero_more_reg}\(.*\)${blank_zero_more_reg}[{]?${blank_zero_more_reg}$"
            ;;
        struct)
            # 结构体定义模式：struct + 名称 + {
            echo "${start_reg}struct${blank_one_more_reg}${keyword_reg}${blank_zero_more_reg}{"
            ;;
        enum_const)
            echo "${start_reg}${keyword_reg}[,]?(${blank_zero_more_reg}|${blank_zero_more_reg}(//|/\*).*)$"
            ;;
        typedef)
            # typedef: typedef struct {} keyword
            echo "${start_reg}typedef${blank_one_more_reg}(struct${blank_zero_more_reg}{.*}|[a-zA-Z_].*${blank_one_more_reg})${blank_zero_more_reg}${keyword_reg}"
            ;;
        all)
            # 通用模式：包含关键词的行
            echo "$keyword"
            ;;
        *)
            echo "$keyword"
            ;;
    esac
}

# 第一步：基础搜索（保留所有结果）
basic_search() {
    local keyword="$1"
    local dirs_str="$2"
    local dirs=()
    IFS=' ' read -ra dirs <<< "$dirs_str"
    local ext_pattern="$3"
    
    echo -e "${YELLOW}=== 第一步：基础搜索 ===${NC}" >&2
    echo -e "关键词: ${CYAN}$keyword${NC}" >&2
    echo -e "目录: ${BLUE}${dirs[*]}${NC}" >&2
    echo -e "文件类型: ${PURPLE}${FILE_EXTENSIONS[*]}${NC}" >&2
    echo -e "文件搜索模式: ${GREEN}${ext_pattern}${NC}" >&2
    echo "" >&2
    
    local all_results_file="./search_results_$$.txt"
    > "$all_results_file"
  
    for dir in "${dirs[@]}"; do
        if [[ !  -d "$dir" ]]; then
            if [[ $VERBOSE -eq 1 ]]; then
                echo -e "${RED}警告: 目录 $dir 不存在，跳过${NC}" >&2
            fi
            continue
        fi
        
        if [[ $VERBOSE -eq 1 ]]; then
            echo -e "${BLUE}搜索目录: $dir${NC}" >&2
        fi
        
        # 使用grep递归搜索
        if [[ -n "$ext_pattern" ]]; then
            find "$dir" -type f | grep "$ext_pattern" | xargs grep -n --color=always  "$keyword" 2>/dev/null || true
        else
            grep -rn --color=always  "$keyword" "$dir" 2>/dev/null || true
        fi
    done > "$all_results_file"
    
    # 显示结果数量
    local result_count=$(grep -c "$keyword" "$all_results_file" 2>/dev/null || echo 0)
    echo -e "${GREEN}基础搜索完成，找到 ${result_count} 个结果${NC}" >&2
    echo "" >&2
    echo -e "============ ${GREEN}搜索结果${NC} ================" >&2
    echo -e "文件名: ${CYAN}$all_results_file${NC}" >&2 
    cat $all_results_file >&2
    echo "" >&2
    echo "$all_results_file"
}

# 第二步：智能筛选
smart_filter() {
    local results_file="$1"
    local keyword="$2"
    local search_type="$3"
 
    echo -e "${YELLOW}=== 第二步：智能筛选 ===${NC}" >&2
    
    local filtered_results_file="./filtered_results_$$.txt"
    local pattern=$(build_search_pattern "$keyword" "$search_type")
    
    if [[ $VERBOSE -eq 1 ]]; then
        echo  "筛选模式: $pattern$" >&2
    fi
    
    # 根据类型进行筛选
    case "$search_type" in
        macro)
            # 筛选宏定义
            grep -E "$pattern" "$results_file" > "$filtered_results_file" 
            ;;
        function)
            # 筛选函数定义（排除声明和调用）
            grep -E "$pattern" "$results_file" | \
            grep -v ";[[:blank:]]*$" > "$filtered_results_file" || true
            ;;
        struct)
            # 筛选结构体定义
            grep -E "$pattern" "$results_file" > "$filtered_results_file" || true
            ;;
        enum_const)
            grep -E "$pattern" "$results_file" > "$filtered_results_file" || true
            ;;
        typedef)
            grep -E "$pattern" "$results_file" > "$filtered_results_file" || true
            ;;
        all)
            # 合并所有类型的结果
            local temp_file="./combined_$$.txt"
            > "$temp_file"  # 清空临时文件
            
            # 依次执行macro、function、struct、enum_const的筛选并去重合并
            local macro_pattern=$(build_search_pattern "$keyword" "macro")
            local function_pattern=$(build_search_pattern "$keyword" "function")
            local struct_pattern=$(build_search_pattern "$keyword" "struct")
            local enum_const_pattern=$(build_search_pattern "$keyword" "enum_const")
            local typedef_pattern=$(build_search_pattern "$keyword" "typedef")
            
            if [[ $VERBOSE -eq 1 ]]; then
                echo "macro筛选模式: $macro_pattern" >&2
                echo "function筛选模式: $function_pattern" >&2
                echo "struct筛选模式: $struct_pattern" >&2
                echo "typedef筛选模式: $typedef_pattern" >&2
            fi
            
            # 筛选macro并追加
            grep -E "$macro_pattern" "$results_file" >> "$temp_file" 2>/dev/null || true
            
            # 筛选function并追加（排除声明）
            grep -E "$function_pattern" "$results_file" | \
            grep -v ";[[:blank:]]*$" >> "$temp_file" 2>/dev/null || true
            
            # 筛选struct并追加
            grep -E "$struct_pattern" "$results_file" >> "$temp_file" 2>/dev/null || true

            grep -E "$enum_const_pattern" "$results_file" >> "$temp_file" 2>/dev/null || true

            grep -E "$typedef_pattern" "$results_file" >> "$temp_file" 2>/dev/null || true

            # 去重并保存结果
            sort -u "$temp_file" > "$filtered_results_file"
            rm -f "$temp_file"
            ;;
    esac
    
    # 显示筛选结果
    local filtered_count=$(wc -l < "$filtered_results_file" 2>/dev/null || echo 0)
    echo -e "${GREEN}筛选完成，剩余 ${filtered_count} 个可能的结果${NC}" >&2
    echo "" >&2
    
    echo -e "============ ${GREEN}筛选结果${NC}=============" >&2
    echo -e "文件名: ${CYAN}$filtered_results_file${NC}" >&2
    cat $filtered_results_file >&2
    echo "" >&2
    echo "$filtered_results_file"
}

# 第三步：结果分析和排序
analyze_results() {
    local results_file="$1"
    local keyword="$2"
    
    echo -e "${YELLOW}=== 第三步：结果分析和排序 ===${NC}"
    
    if [[ ! -s "$results_file" ]]; then
        echo -e "${RED}没有找到匹配的结果${NC}"
        return
    fi
    
    echo -e "${CYAN}最可能的定义位置：${NC}"
    echo "----------------------------------------"
    
    # 使用数组存储文件信息
    local files=()
    local scores=()
    local lines_content=()
    
    # 处理结果文件
    while IFS=: read -r file line content; do
        [[ -z "$file" || -z "$line" ]] && continue
        # 评分规则
        local score=0

        echo "$file" | grep -q -E "(default)" && ((score += 5))
        echo "$file" | grep -q -E "(linux/|genric)" && ((score += 10))
        
        arch_pattern=$(IFS='|'; echo "${ARCH[*]}")
        echo "$file" | grep -q -E "($arch_pattern)" && ((score += 20))
        # 查找文件是否已存在
        local found=0
        for i in "${!files[@]}"; do
            if [[ "${files[i]}" == "$file" ]]; then
                scores[i]=$((scores[i] + score))
                lines_content[i]="${lines_content[i]}"$'\n'"$line:$content"
                found=1
                break
            fi
        done
        
        if [[ $found -eq 0 ]]; then
            if [[ ${#files[@]} -eq 1 && "${files[0]}" == "()" ]]; then
                files[0]="$file"
            else
                files[${#files[@]}]="$file"
            fi

            if [[ ${#scores[@]} -eq 1 && "${scores[0]}" == "()" ]]; then
                scores[0]="$score"
            else
                scores[${#scores[@]}]="$score"
            fi

            if [[ ${#lines_content[@]} -eq 1 && "${lines_content[0]}" == "()" ]]; then
                lines_content[0]="$line:$content"
            else
                lines_content[${#lines_content[@]}]="$line:$content"
            fi
        fi
    done < "$results_file"

    # 直接排序并输出，避免中间数组
    for ((i=0; i<${#files[@]}; i++)); do
        echo "${scores[i]}:${files[i]}:$i"
    done | sort -t: -k1,1nr | while IFS=: read -r score file index; do
        echo -e "${BLUE}文件: $file (置信度: $score)${NC}"
        echo -e "${BLUE}匹配行:${NC}"
    
        # 处理该文件的所有匹配行
        local content_block="${lines_content[index]}"
        while IFS=: read -r line content; do
            [[ -n "$line" && -n "$content" ]] && echo -e "  行 $line: $content"
        done <<< "$content_block"

        # 只对第一个文件显示第一个匹配行的上下文
        if [[ $first_file_processed -eq 0 ]]; then
            # 获取第一个匹配行号
            first_match_line=$(echo "$content_block" | head -1 | cut -d: -f1)
            if [[ -n "$first_match_line" && "$first_match_line" =~ ^[0-9]+$ && -f "$file" ]]; then
                echo -e "${YELLOW}=== 第一个匹配行的上下文（前后各${NUM}行）===${NC}"
                
                start_line=$((first_match_line - ${NUM}))
                [[ $start_line -lt 1 ]] && start_line=1
                end_line=$((first_match_line + ${NUM}))
                
                sed -n "${start_line},${end_line}p" "$file" | while IFS= read -r context_line; do
                    current_line=$((start_line + line_count))
                    if [[ $current_line -eq $first_match_line ]]; then
                        printf "  ${RED}- %d: %s${NC}\n" "$current_line" "$context_line"
                    else
                        echo -e "    $current_line: $context_line"
                    fi
                    ((line_count++))
                done
            fi
            first_file_processed=1
        fi
 
        echo "----------------------------------------"
    done

}

# 第四步：保存详细结果
save_detailed_results() {
    local basic_file="$1"
    local filtered_file="$2"
    local keyword="$3"
    
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local output_file="search_results_${keyword}_${timestamp}.txt"
    
    echo -e "${YELLOW}=== 第四步：保存结果 ===${NC}" >&2
    
    {
        echo "搜索时间: $(date)"
        echo "关键词: $keyword"
        echo "搜索类型: $SEARCH_TYPE"
        echo "搜索目录: ${SEARCH_DIRS[*]}"
        echo "文件类型: ${FILE_EXTENSIONS[*]}"
        echo ""
        if [[ -f "$basic_file" ]]; then
            echo "=== 基础搜索结果 ==="
            local filtered_count=$(wc -l < "$basic_file" 2>/dev/null || echo 0)
            echo -e "${GREEN}基础搜索完成, 有 ${filtered_count} 个可能的结果${NC}"
            cat "$basic_file"
        else
            echo "=== 基础搜索结果 ==="
            echo "无结果"
        fi
        echo ""
        if [[ -f "$filtered_file" ]]; then
            echo "=== 筛选后结果（${keyword}）==="
            local filtered_count=$(wc -l < "$filtered_file" 2>/dev/null || echo 0)
            echo -e "${GREEN}筛选完成, 剩余 ${filtered_count} 个可能的结果${NC}"
            cat "$filtered_file"
        else
            echo "=== 筛选后结果（${keyword}）==="
            echo "无结果"
        fi
    } > "$output_file"
    
    if [[ $SAVE -eq 0 ]]; then
        echo -e "${GREEN}选择不保存输出结果！${NC}" >&2
        rm -rf "$output_file"
        echo -e "${GREEN}已删除输出结果！${NC}" >&2
        return
    fi

    echo -e "${GREEN}详细结果已保存到: ${CYAN}$output_file${NC}" >&2
    echo "$output_file"
}

# 主函数
main() {
    echo -e "${PURPLE}开始搜索 Linux 内核定义...${NC}"
    echo ""
    
    # 第一步：基础搜索
    local ext_pattern=$(build_extension_pattern)
    local basic_results=$(basic_search "$KEYWORD" "${SEARCH_DIRS[*]}" "$ext_pattern")

    # 第二步：智能筛选
    local filtered_results=$(smart_filter "$basic_results" "$KEYWORD" "$SEARCH_TYPE")
    
    # 第三步：结果分析
    analyze_results "$filtered_results" "$KEYWORD"
    
    # 第四步：保存结果
    local output_file=$(save_detailed_results "$basic_results" "$filtered_results" "$KEYWORD")
    
    # 清理临时文件
    rm -f "$basic_results" "$filtered_results"
    
    echo ""
    echo -e "${GREEN}搜索完成！${NC}"
    echo -e "支持定义类型: macro function struct enum_const typedef all"
}

# 执行主函数
main "$@"